### Two-Stage Qualitative Tradeoff Resolution (2S-QTR) function library
## Main model: Section 4 in paper

MaxLikeSR_2sQTR_LOO_crossVal = function(data, sub_sample, init){
  # Leave-one-out Cross Validation (LOO CV)
  start <- Sys.time()  
  data = sample_SR(data,sub_sample)
  index = data$studentID                                                            
  data = split(data,f = index)  
  
  dfs_LOOCV = lapply(data, LOOCV_train_test_split)
  results = lapply(dfs_LOOCV, LOO_crossVal_train_predict_2sQTR, init)
  
  epsilon_coef = unlist(lapply(results, {function(x) mean(x[["epsilon.coef"]])}), use.names = FALSE)
  epsilon_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["epsilon.coef"]]))}), use.names = FALSE)
  
  alpha_coef = unlist(lapply(results, {function(x) mean(x[["alpha.coef"]])}), use.names = FALSE)
  alpha_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["alpha.coef"]]))}), use.names = FALSE)
  
  beta_coef = unlist(lapply(results, {function(x) mean(x[["beta.coef"]])}), use.names = FALSE)
  beta_coef_std = unlist(lapply(results, {function(x) sqrt(var(x[["beta.coef"]]))}), use.names = FALSE)
  
  score = unlist(lapply(results, {function(x) sum(x$accuracy)}), use.names = FALSE)
  results_scores = as.data.frame(cbind(studentID = unique(index), score = score,
                                       epsilon_mean = epsilon_coef, alpha_mean = alpha_coef, beta_mean = beta_coef,
                                       epsilon_std = epsilon_coef_std, alpha_std = alpha_coef_std, beta_std = beta_coef_std))
  
  print("Two-Stage Qualitative Tradeoff Resolution - Leave-One-Out Test Accuracy Summary:")
  print(paste("Average:", mean(results_scores$score)))
  print(paste("Standard error:", sqrt(var(results_scores$score)/length(results_scores$score))))
  
  end = Sys.time()                                                                  
  timer = difftime(end,start,units = "secs")
  print(paste("Time taken for MLE algorithim and Cross-Validation:",seconds_to_period(round(timer[[1]],2))))
  return(list(estimations = results, summary = results_scores))
}

LOOCV_train_test_split = function(data){
  result = list()
  for(i in 1:length(data$qnum)){
    test_df = data[i, ]
    train_df = data[-i, ]
    result[[i]] <- list(train = train_df, test = test_df)
  }
  result
}


LOO_crossVal_train_predict_2sQTR = function(LOO_train_test_folds, init){
  results = bind_rows(lapply(LOO_train_test_folds, LOO_train_test_2sQTR, init))
  results
}


LOO_train_test_2sQTR = function(train_test_dfs, init){
  train_df = train_test_dfs[["train"]]
  test_df = train_test_dfs[["test"]]
  
  results = wrapMLE_2sQTR(train_df, init, log_lik_2sQTR)
  results_summary = stat_sum_2sQTR(results)
  results_summary$studentID = test_df$studentID
  
  param = list(epsilon = results_summary$epsilon.coef, alpha = results_summary$alpha.coef, beta = results_summary$beta.coef)
  test_predict_results = LOO_test_predict_2sQTR(test_df, param)
  results_summary = cbind(results_summary, test_qnum = test_df$qnum, test_choice = test_df$choice, test_predict_results)
  results_summary
}


LOO_test_predict_2sQTR = function(test_df, param){
  l1 = (1-param[["epsilon"]])*param[["alpha"]]*param[["beta"]] + (param[["epsilon"]]/5)
  l2 = (1-param[["epsilon"]])*param[["alpha"]]*(1-param[["beta"]]) + (param[["epsilon"]]/5)
  l3 = param[["epsilon"]]/5
  l4 = (1-param[["epsilon"]])*(1-param[["alpha"]])*(1-param[["beta"]]) + (param[["epsilon"]]/5)
  l5 = (1-param[["epsilon"]])*(1-param[["alpha"]])*(param[["beta"]]) + (param[["epsilon"]]/5)
  
  lottery_prob = c(l1,l2,l3,l4,l5)
  predicted_choice = unname(which.max(lottery_prob))
  accuracy = if_else(predicted_choice == test_df$choice, 1, 0)
  list(pred_choice = predicted_choice, accuracy = accuracy)
}


wrapMLE_2sQTR = function(data,init,logLikFun){
  QTR_MaxLike_res = mle2(minuslogl = logLikFun, data = list(data=data), start = init,
                                  lower = c(epsilon=0.01, alpha=0.01, beta=0.01),
                                  upper = c(epsilon=1, alpha=1, beta=1))
  QTR_MaxLike_res
}


stat_sum_2sQTR = function(stat){
  results = data.frame("studentID"=NA, "epsilon.coef"=NA, "alpha.coef" = NA, "beta.coef" = NA,
                       "std.epsilon" = NA, "std.alpha" = NA, "std.beta" = NA, 
                       "p.val.epsilon"=NA, "p.val.alpha"=NA, "p.val.beta"=NA,
                       "logLikelihood"=NA,"AIC"=NA)
  
  summary = summary(stat)
  results["studentID"]       = NA
  results["epsilon.coef"]    = summary@coef[1,1]
  results["alpha.coef"]      = summary@coef[2,1]
  results["beta.coef"]       = summary@coef[3,1]
  results["std.epsilon"]     = summary@coef[1,2]
  results["std.alpha"]       = summary@coef[2,2]
  results["std.beta"]        = summary@coef[3,2]
  results["p.val.epsilon"]   = summary@coef[1,4]
  results["p.val.alpha"]     = summary@coef[2,4]
  results["p.val.beta"]      = summary@coef[3,4]
  results["logLikelihood"]   = logLik(stat)
  results["AIC"]             = AIC(stat)
  results
}


calc_2sQTR = function(epsilon, alpha, beta, choice, data){
  if(choice==1){(1-epsilon)*alpha*beta + (epsilon/5)}
  else if(choice==2){(1-epsilon)*alpha*(1-beta) + (epsilon/5)}
  else if(choice==3){epsilon/5}
  else if(choice==4){(1-epsilon)*(1-alpha)*(1-beta) + (epsilon/5)}
  else{(1-epsilon)*(1-alpha)*beta + (epsilon/5)}
}



log_lik_2sQTR <- function(epsilon, alpha, beta, data){
  # calculate likelihood of a choice
  Calc = function(data){
    choice = data[["choice"]]
    calc = calc_2sQTR(epsilon, alpha, beta, choice, data)
    calc
  }
  
  
  Res = apply(data, 1, Calc)
  
  #calculating log likelihood
  log_lik_i <-  log(Res)
  
  -sum(unlist(log_lik_i))
}


sample_SR = function(data,sub_sample){
  if(sub_sample == "full"){data = data}
  if(sub_sample == "part1"){data = data %>% filter(qnum<19)}
  if(sub_sample == "part2"){data = data %>%  filter(qnum>18)}
  return(data)
}






